//065959264	Limor Bagizada	040199770	Sarit Chicotay

package Models;

import java.util.ArrayList;
import java.util.Set;
import java.util.TreeMap;

import Base.NGram;

/**
 * 
 * Class implementing the BackOff discounting model. 
 * The model is implemented using Lidstone smoothing for the discounted probabilities.
 *
 */
public class BackOffModel extends BaseModel{

	// discounting model used for probabilities 
	private LidstoneModel discountModel;
	
	// discounting model used for back-off probabilities 
	private LidstoneModel backOffModel;
	
	// for improving performance - calculate alpha for each conditioning word only once
	private TreeMap<String, Double> alpha;
	
	// constructor
	public BackOffModel(ArrayList<String> content) throws Exception {
		
		// initialize model with first 90% of words as training, the rest will be used for test
		super(content, 0.9, 2);
	
		discountModel = new LidstoneModel(content, 2);//bigram
		backOffModel = new LidstoneModel(content, 1);//unigram
		alpha = new TreeMap<String, Double>();
	}

	@Override
	/**.
	 * P(x) - calculate the probability of an event
	 */
	public double CalculateProbability(NGram event) {
		double result = 0;
		
		if (event.getNGram() == 1 || training.getCount(event)>0) // calculate probability using Lidstone discounting
		{
			result = discountModel.CalculateProbability(event);
		}
		else // event count equals zero - use back-off
		{
			double backOffProbability = backOffModel.CalculateProbability(new NGram(event.getWord(),null));
			double backOffAlpha = calculate_alpha(event.getConditioningWord());
			result = (double) backOffAlpha * backOffProbability;
		}
		return result;
	}

	
	/**
	 * P(x) - calculate the probability of an event that appears r times in the corpus using lidstone discounting model
	 */
	public double CalculateProbability(int r, String conditioningWord) throws Exception {
		double backOffProbability = backOffModel.CalculateProbability(r);
		double backOffAlpha = calculate_alpha(conditioningWord);
		double result = (double) backOffAlpha * backOffProbability;
		return result;
	}
	
	// calculate the multiplier of 
	private double calculate_alpha(String w){
		
		if (alpha.containsKey(w))
			return alpha.get(w);
		
		double tempAlpha = 0;
		double beta=0;
		double sum_counter=0;
		double sum_denominator=0;
		// loop all existing words
		for (String word : training.getUnigramVocubularyKeySet())
		{
			NGram event = new NGram(word, w);
			if (training.getCount(event)>0) // only for existing events 
			{
				sum_counter+=discountModel.CalculateProbability(event);
				sum_denominator+=backOffModel.CalculateProbability(new NGram(word, null));
			}
		}
		beta=(double)1-sum_counter;
		tempAlpha = (double)(beta/((double)1-sum_denominator));
		alpha.put(w, tempAlpha);
		return tempAlpha;
	}
	
	public Set<String> getUnigramVocubularyKeySet()
	{
		return this.training.getUnigramVocubularyKeySet();
	}
	
	@Override
	public double debugging() throws Exception {
				
		// for each conditional word - total probabilities should be equal 1
		for (String conditioningWord : training.getUnigramVocubularyKeySet())
		{
			double result=0;
			double temp=0;
			
			for (String word : training.getUnigramVocubularyKeySet())
			{
				temp = CalculateProbability(new NGram(word, conditioningWord));
				result += temp;
			}
		
			result += debugging_N0(conditioningWord);
			
			System.out.println("debugging " + conditioningWord + " :" + result);
		}
		
		return 0;
	}
	
	// for words that are not in the corpus
	public double debugging_N0(String conditioningWord) throws Exception
	{
		double result = 0;
		int N_0 = this.return_N_r(0);
		if (N_0 > 0)
		{
			// get some word that was not in the training
			double temp  = CalculateProbability(0, conditioningWord);
			 
			 // multiply the probability of such a word with N0
		    result+= (double)(temp*(double)N_0);
		}
		return result;
	}

	@Override
	public double CalculateProbability(int r) throws Exception {
		throw new Exception("Cannot use back-off CalculateProbability for count r");
	}
	
}